from datetime import datetime
from typing import Union, Optional

from sqlalchemy import update, CursorResult, text, func, insert, delete
from sqlalchemy.ext.asyncio import AsyncSession
from sqlmodel import select, desc, col, asc, and_
from sqlmodel.sql._expression_select_cls import SelectOfScalar, Select

from src.base.beans import PageResult
from src.base.constants import id_generator
from src.domain.sys_dept import SysDept, SysDeptReq


async def get_dept_by_id(dbSession: AsyncSession, id: int) -> SysDept:
    result = await dbSession.execute(select(SysDept).where(SysDept.dept_id == id).where(SysDept.del_flag == False))
    return result.scalar_one_or_none()


async def get_dept_list_by_ids(dbSession: AsyncSession, ids: list[int]) -> list[SysDept]:
    result = await dbSession.execute(select(SysDept).where(col(SysDept.dept_id).in_(ids)))
    return list(result.scalars().all())


def where_handle(stmt: Union[Select, SelectOfScalar], entity: SysDeptReq = None):
    stmt = stmt.where(SysDept.del_flag == False)
    if entity:
        if entity.dept_id:
            stmt = stmt.where(SysDept.dept_id == entity.dept_id)
        if entity.dept_name:
            stmt = stmt.where(col(SysDept.dept_name).like(f"%{entity.dept_name}%"))
        if entity.status is not None:
            stmt = stmt.where(SysDept.status == entity.status)
        if entity.create_time_range:
            start_time = datetime.strptime(entity.create_time_range[0], "%Y-%m-%d").date()
            end_time = datetime.strptime(entity.create_time_range[1], "%Y-%m-%d").date()
            stmt = stmt.where(
                and_(func.date(SysDept.create_time) >= start_time, func.date(SysDept.create_time) <= end_time))
    return stmt


async def get_dept_list_by_condition(dbSession: AsyncSession, entity: Optional[SysDeptReq] = None) -> list[SysDept]:
    stmt = select(SysDept)
    stmt = where_handle(stmt, entity)
    stmt = stmt.order_by(desc(SysDept.sort))
    result = await dbSession.execute(stmt)
    return list(result.scalars())


async def get_dept_page_by_condition(dbSession: AsyncSession, entity: Optional[SysDeptReq] = None) -> PageResult:
    list_stmt = select(SysDept)
    list_stmt = where_handle(list_stmt, entity)
    if entity.sorter:
        list_stmt = list_stmt.order_by(text(entity.sorter))
    else:
        list_stmt = list_stmt.order_by(asc(SysDept.sort))
    list_stmt = list_stmt.limit(entity.pageSize).offset(entity.pageSize * (entity.current - 1))
    data = await dbSession.execute(list_stmt)
    data = list(data.scalars())

    total_stmt = select(func.count(SysDept.dept_id))
    total_stmt = where_handle(total_stmt, entity)
    total_result = await dbSession.execute(total_stmt)
    total = total_result.scalar_one_or_none()
    return PageResult(total=total, data=data, current=entity.current, pageSize=entity.pageSize)


async def insert_dept(dbSession: AsyncSession, entity: SysDeptReq):
    entity.dept_id = next(id_generator)
    insert_data = SysDept(**entity.model_dump()).model_dump(exclude_unset=True, exclude_none=True)
    result = await dbSession.execute(insert(SysDept).values(**insert_data))
    # entity.dept_id = result.lastrowid
    return int(result.rowcount)


async def update_dept(dbSession: AsyncSession, entity: SysDeptReq) -> int:
    update_data = SysDept(**entity.model_dump()).model_dump(exclude_unset=True, exclude_none=True)
    stmt = update(SysDept).values(**update_data).where(col(SysDept.dept_id).__eq__(entity.dept_id))
    result: CursorResult = await dbSession.execute(stmt)
    return int(result.rowcount)


async def delete_dept(dbSession: AsyncSession, id: int = None) -> int:
    # stmt = update(SysDept).values({"del_flag": True}).where(col(SysDept.dept_id).__eq__(id))
    stmt = delete(SysDept).where(col(SysDept.dept_id).__eq__(id))
    result: CursorResult = await dbSession.execute(stmt)
    return int(result.rowcount)


async def delete_depts(dbSession: AsyncSession, ids: list[int] = None) -> int:
    # stmt = update(SysDept).values({"del_flag": True}).where(col(SysDept.dept_id).in_(ids))
    stmt = delete(SysDept).where(col(SysDept.dept_id).in_(ids))
    result: CursorResult = await dbSession.execute(stmt)
    return int(result.rowcount)


async def get_dept_with_childes_by_ids(dbSession: AsyncSession, dept_ids: list[int]) -> list[SysDept]:
    menu_sql = text(f"WITH RECURSIVE dept_tree AS ("
                    f"SELECT dept_id, parent_id, dept_name FROM sys_dept WHERE dept_id in :dept_ids "
                    f"UNION ALL "
                    f"SELECT d.dept_id, d.parent_id, d.dept_name  FROM sys_dept d INNER JOIN dept_tree t ON d.parent_id = t.dept_id)"
                    f"SELECT * FROM dept_tree")
    menu_stmt = select(SysDept).from_statement(menu_sql)
    result = await dbSession.execute(menu_stmt, params={"dept_ids": tuple(dept_ids)})
    return list(result.scalars())
